import cv2
import numpy as np
import os
import json


class ImCutMergeByWindow(object):
    def __init__(self):
        self.__main_img = None  # 合并图像时用到的空白大图
        self.__temp_dir = "./temp_result_dir"  # 默认的临时文件夹，默认将生成的切片存入其中
        self.__img_info_file = "info.txt"  # 图像信息数据存于此处

    def __creat_dir(self, result):
        """
        若目录不存在则创建目录
        :param result: 要创建的目录的名字
        """
        if not os.path.exists(result):
            os.mkdir(result)

    def __write_basic_info_to_file(self, img, clip_w_pixels, clip_h_pixels, step_w, step_h, total_clip_num, clip_type):
        """
        这个函数用于将分割后图像的数据写入info.txt文件中去，方便合并时从中读取对应的所需参数
        :param img: 原始图像，其本身不写入，此处的传入是为了得到原始图像的基本信息
        :param clip_w_pixels: 切片宽度像素数目
        :param clip_h_pixels: 切片高度像素数目
        :param step_w: 横向步长
        :param step_h: 纵向步长
        :param total_clip_num: 最后生成的总的切片数目
        :param clip_type: 切片存储的类型
        """

        # 建立对应信息的字典，用于生成对应的json字符串
        img_info = {"img_h": img.shape[0], "img_w": img.shape[1], "img_c": img.shape[2],
                    "clip_w_pixels": clip_w_pixels, "clip_h_pixels": clip_h_pixels, "step_w": step_w,
                    "step_h": step_h, "total_clip_num": total_clip_num, "clip_type": clip_type}
        # 打开文件info.txt
        info_file = open("%s\\%s" % (self.__temp_dir, self.__img_info_file), 'w')
        # 生成json
        temp_content = json.dumps(img_info)
        # 计算json的字节长度，并首先第一行写入
        basic_info_len = len(temp_content.encode())
        info_file.write("%s\n" % basic_info_len)
        # 写入json到文件
        info_file.write(temp_content)
        info_file.flush()
        info_file.close()

    def __read_basic_info_from_file(self, src_dir):
        """
        与写入数据到文件对应的从info.txt中读取数据
        :param src_dir: 目标文件夹，必须包含info.txt
        :return: 包含图像合并所需基本信息的字典
        """
        if not os.path.exists(src_dir):
            return None
        temp_file = open("%s\\%s" % (src_dir, self.__img_info_file), 'r')
        basic_info_len_str = temp_file.readline()
        # 得到第一行的json长度，不要忽略了还有’\n‘
        step1 = len(basic_info_len_str.encode()) + 1
        basic_info_len = int(basic_info_len_str)
        # 跳过首行
        temp_file.seek(step1, 0)
        # 读取对应长度的信息
        content = temp_file.read(basic_info_len)
        img_info = json.loads(content)
        temp_file.flush()
        temp_file.close()
        return img_info

    def __generate_file_name(self, end_with, *args):
        """
        根据传入的参数生成文件名称字符串 形如 a_b_c_d.bmp
        :param args:组成文件名称的参数
        :param end_with:文件后缀
        :return:完整的文件名字
        """
        name_str = str(args[0])
        for i in range(1, len(args)):
            name_str = name_str + "_" + str(args[i])
        name_str = name_str + end_with
        return name_str

    def __decode_file_name(self, file_name_str):
        """
        将传入的文件名拆分如将 a_b_c_d.jpg 拆分为  a  b  c  d
        :param file_name_str: 文件名字符串
        :return: 参数元组
        """
        args = file_name_str.split('_')
        result_list = []
        for i in range(0, len(args) - 1):
            result_list.append(int(args[i]))
        temp_s = args[len(args) - 1].split('.')
        result_list.append(int(temp_s[0]))
        return tuple(result_list)

    def __generate_zeros_img(self, *shape):
        """
        生成一张黑色图像，并保存在self.__main_img中
        :param shape: 要生成图像的信息
        :return: 创建是否成功
        """
        del self.__main_img
        self.__main_img = np.zeros(list(shape), dtype=np.uint8)  # 生成一张空白黑色图像
        if self.__main_img is not None:
            return True
        else:
            return False

    def img_cut(self, img, clip_w_pixels, clip_h_pixels, step_w, step_h, result_dir=None, clip_file_type=".bmp"):
        """
        将是原始图像按照给定参数分割成单张小图像，并将其保存至默认目录
        :param clip_file_type: 切片后的小图像的存储类型 默认为bmp格式
        :param img: 原始图像
        :param clip_w_pixels: 单张切片的宽度像素数目
        :param clip_h_pixels: 单张切片的高度像素数目
        :param step_w: 滑动窗口的横向步长
        :param step_h: 滑动窗口的纵向步长
        :param result_dir: 分割后的结果的存放目录，为None则使用默认值
        """
        # img[clip_top:clip_bottom, clip_left:clip_right]
        if img is not None:
            # 根据传入的参数是否为空来选择是使用默认路径还是外部指定的路径
            if result_dir is not None:
                self.__creat_dir(result_dir)  # 建立临时文件夹
            else:
                self.__creat_dir(self.__temp_dir)
            # 初始化临时参数
            img_h, img_w, clip_top, clip_bottom, clip_left, clip_right = img.shape[0], img.shape[
                1], 0, clip_h_pixels, 0, clip_w_pixels
            total_clips = 0  # 用于统计最后一共切割获得了多少张图像
            while True:
                clip_left, clip_right = 0, clip_w_pixels  # 本次切片的左边界，右边界
                while True:
                    # 生成切片临时图像，并且以切片位置信息为参数生成图像的名称写入外存
                    temp_img = img[clip_top:clip_bottom, clip_left:clip_right]
                    # 生成本张切片的文件名
                    img_name = self.__generate_file_name(clip_file_type, clip_top, clip_bottom, clip_left, clip_right)
                    # 写入对应文件夹中
                    cv2.imwrite("%s\\%s" % (self.__temp_dir, img_name), temp_img)
                    total_clips = total_clips + 1  # 切片总数加一
                    if clip_right == img_w:
                        break  # 当到达图像右边界时结束循环
                    clip_left = clip_left + step_w  # 更新切片左右边界
                    clip_right = clip_right + step_w
                    # 防止右边界超出原图像的大小范围，即切片的最后一张的大小可能会小于clip_w_pixels要求的大小
                    if clip_right > img_w:
                        clip_right = img_w
                if clip_bottom == img_h:
                    break  # 当到达图像底部时结束循环
                clip_top = clip_top + step_h  # 更新切片上下边界
                clip_bottom = clip_bottom + step_h
                # 与左右边界同理，防止下边界超出原图像的大小范围
                if clip_bottom > img_h:
                    clip_bottom = img_h
            # 在所有切片切割完成时，将全部信息写入info.txt中
            self.__write_basic_info_to_file(img, clip_w_pixels,
                                            clip_h_pixels, step_w,
                                            step_h, total_clips, clip_file_type)

    def img_merge(self, src_dir=None):
        """
        将指定文件夹或者默认文件夹中的图像合并为一张大图
        :param src_dir: 小图像的路径
        :return: 合并后的大图
        """
        target_src = None
        if src_dir is not None:
            target_src = src_dir
        else:
            target_src = self.__temp_dir
        img_basic_info = self.__read_basic_info_from_file(target_src)
        if img_basic_info is None:
            print("src_dir路径无效")
            return None
        else:
            """
            切割信息写入时写入格式如下
             {"img_h": img.shape[0], "img_w": img.shape[1], "img_c": img.shape[2],
                    "clip_w_pixels": clip_w_pixels, "clip_h_pixels": clip_h_pixels, "step_w": step_w,
                    "step_h": step_h, "total_clip_num": total_clip_num, "clip_type": clip_type}
            """

            # 原图的 高度，宽度，通道数
            img_h = img_basic_info["img_h"]
            img_w = img_basic_info["img_w"]
            img_c = img_basic_info["img_c"]
            # clip_w_pixels = img_basic_info["clip_w_pixels"]
            # clip_h_pixels = img_basic_info["clip_h_pixels"]
            # step_w = img_basic_info["step_w"]
            # step_h = img_basic_info["step_h"]
            # 文件夹中应该包含的切片数目
            total_clip_num = img_basic_info["total_clip_num"]
            # 切片的文件类型
            clip_file_type = img_basic_info["clip_type"]
            # 得到对应文件夹中对应类型文件的列表
            img_list = self.__get_file_list(target_src, total_clip_num, clip_file_type)
            # 遍历上面得到的文件列表，并将切片合并到大图中
            if img_list is not None:
                if self.__generate_zeros_img(img_h, img_w, img_c):
                    for img_name in img_list:
                        # 根据文件的名字来确定本张图片在原图中的位置，此处通过解析文件名字获得一个范围
                        top, bottom, left, right = self.__decode_file_name(img_name)
                        # 读取图像的数据
                        temp_img = cv2.imread("%s\\%s" % (target_src, img_name))
                        # 将数据覆盖到准备好的空白矩阵中
                        self.__main_img[top:bottom, left:right] = temp_img
            return self.__main_img

    def __get_file_list(self, tmp_dir, check_size, require_file_type):
        """
        获取文件中指定类型文件的文件名列表
        :param tmp_dir: 指定文件夹
        :param check_size: 文件的数量
        :param require_file_type:
        :return:
        """
        result_list = []
        if os.path.exists(tmp_dir):
            for i in os.listdir(tmp_dir):
                if i.endswith(require_file_type):
                    result_list.append(i)
            # print(len(temp_file_list))
            if len(result_list) < check_size:
                print("图片数量不足")
                return None
            else:
                return result_list
        return None


# 此处 生成一个对象，外部使用时只需要导入此对象就行
imc = ImCutMergeByWindow()
